import pandas as pd
from transformers import BertTokenizer
from torch.utils.data import Dataset, DataLoader
import torch

# 定义数据处理函数
class QNLIDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_len=128, num_samples=None):
        self.data = pd.read_csv(data_path, sep="\t", on_bad_lines='skip')
        if num_samples is not None:
            self.data = self.data.iloc[:num_samples]
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        question = self.data.iloc[idx, 1]  # 问题在第2列
        sentence = self.data.iloc[idx, 2]  # 句子在第3列
        label = self.data.iloc[idx, 3]  # 标签在第4列

        # 将标签转换为数值
        label_map = {"entailment": 0, "not_entailment": 1}
        label = label_map[label]

        # 构造BERT输入
        encoding = self.tokenizer.encode_plus(
            question,
            sentence,
            add_special_tokens=True,
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            return_tensors="pt",
        )

        return {
            "input_ids": encoding["input_ids"].flatten(),
            "attention_mask": encoding["attention_mask"].flatten(),
            "labels": torch.tensor(label, dtype=torch.long),
        }
    
